from transformers import BertModel  # 导入 Hugging Face 提供的 BERT 模型类
import torch  # 导入 PyTorch 库，用于构建深度学习模型

# 定义设备信息（如果 GPU 可用，则使用 GPU；否则使用 CPU）
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(DEVICE)  # 打印当前使用的设备信息（GPU 或 CPU）

# 加载预训练的 BERT 模型，并将其移动到指定设备上
pretrained = BertModel.from_pretrained(
    r"D:\SentimentBERT\model\bert-base-chinese\models--bert-base-chinese\snapshots\c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f").to(
    DEVICE)
print(pretrained)  # 打印加载的预训练模型信息，确认模型加载成功


# 定义下游任务模型（增量模型）
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()  # 调用父类构造函数，初始化模型
        # 设计全连接网络，实现二分类任务
        self.fc = torch.nn.Linear(768, 2)  # 定义一个线性层，输入维度为 768（BERT 的隐藏层维度），输出维度为 2（二分类任务）

    def forward(self, input_ids, attention_mask, token_type_ids):
        """
        前向传播函数，定义模型的计算逻辑
        :param input_ids: 输入文本的 token ID
        :param attention_mask: attention mask（用于忽略填充部分）
        :param token_type_ids: token 类型 ID（用于区分句子 A 和 B）
        :return: 模型的输出 logits
        """
        # 冻结 BERT 模型的参数，使其不参与训练
        with torch.no_grad():  # 禁用梯度计算，避免更新 BERT 的参数
            out = pretrained(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
            # 使用预训练的 BERT 模型提取特征，out 包含最后一层的隐藏状态

        # 增量模型参与训练
        out = self.fc(out.last_hidden_state[:, 0])
        # 取 BERT 输出的最后一层隐藏状态的第一个 token（[CLS] token）作为句子的表示
        # 将其输入到全连接层中，得到最终的分类结果
        return out  # 返回模型的输出 logits
    # 逻辑的实现
    # 1.加载预训练模型
    # 2.自定义一个继承自 torch.nn.Module 的模型类。
    # 添加一个全连接层，用于完成二分类任务。
    # 3.前向传播逻辑
    # 冻结 BERT 模型的参数，只训练下游任务的全连接层。
    # 使用 [CLS] token 的隐藏状态作为句子的表示，并输入到全连接层中